import csv
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from scipy.stats import norm

# 气体参数（edep单位为MeV，需转换为eV：1 MeV = 1e6 eV）
Wi = 20 # 加权平均W值（eV/电荷）
Fano = 0.3  # Fano因子

def read_deposits_info(csv_path):
    deposits = []
    with open(csv_path, 'r') as f:
        reader = csv.reader(f)
        in_deposits = False
        header = None
        for row in reader:
            row = [cell.strip() for cell in row]
            if not row:
                continue
            if "Deposits Info" in row:
                in_deposits = True
                try:
                    header = next(reader)
                except StopIteration:
                    print("错误：存款信息部分缺少表头")
                    return []
                continue
            if in_deposits and header:
                try:
                    # 提取必要字段（edep单位为MeV）
                    pre_x = float(row[header.index('pre_pos_x')])
                    pre_y = float(row[header.index('pre_pos_y')])
                    pre_z = float(row[header.index('pre_pos_z')])
                    post_x = float(row[header.index('post_pos_x')])
                    post_y = float(row[header.index('post_pos_y')])
                    post_z = float(row[header.index('post_pos_z')])
                    edep_MeV = float(row[header.index('edep')])  # 单位为MeV
                    
                    deposits.append({
                        'pre_pos_x': pre_x,
                        'pre_pos_y': pre_y,
                        'pre_pos_z': pre_z,
                        'post_pos_x': post_x,
                        'post_pos_y': post_y,
                        'post_pos_z': post_z,
                        'edep_MeV': edep_MeV
                    })
                except ValueError as e:
                    print(f"数据解析错误：{e}")
    return deposits

def process_deposits(deposits):
    charge_positions = []
    for idx, dep in enumerate(deposits):
        # 计算中点位置
        mid_x = (dep['pre_pos_x'] + dep['post_pos_x']) / 2
        mid_y = (dep['pre_pos_y'] + dep['post_pos_y']) / 2
        mid_z = (dep['pre_pos_z'] + dep['post_pos_z']) / 2
        
        edep_MeV = dep['edep_MeV']
        if edep_MeV <= 0:
            print(f"能量沉积 {idx+1}: edep={edep_MeV} MeV（≤0），跳过")
            continue
        
        edep_eV = edep_MeV * 1e6  # MeV转换为eV
        m_ion = edep_eV / Wi  # 平均电离电荷数
        s_ion = np.sqrt(Fano * m_ion)  # Fano噪声标准差
        
        # 生成符合正态分布的电荷数（允许0，但打印详细信息）
        n_ion_float = norm.rvs(loc=m_ion, scale=s_ion)
        n_ion = max(0, int(n_ion_float))
        
        print(f"能量沉积 {idx+1}: edep={edep_MeV:.6f} MeV → {edep_eV:.2f} eV → m_ion={m_ion:.2f} → n_ion={n_ion}")
        
        # 添加扩散（Dt=0.01mm, Dl=0.01mm）
        Dt = 0.008
        Dl = 0.008
        cov = [[Dt**2, 0, 0], [0, Dt**2, 0], [0, 0, Dl**2]]
        for _ in range(n_ion):
            noise = np.random.multivariate_normal([0, 0, 0], cov)
            noisy_pos = (
                mid_x + noise[0],
                mid_y + noise[1],
                mid_z + noise[2]
            )
            charge_positions.append(noisy_pos)
    
    return np.array(charge_positions)

def plot_2d_projection(ax, x_data, y_data, title, xlabel, ylabel):
    if len(x_data) == 0 or len(y_data) == 0:
        return
    x_min, x_max = np.min(x_data), np.max(x_data)
    y_min, y_max = np.min(y_data), np.max(y_data)
    # 添加10%的边距
    x_pad = (x_max - x_min) * 0.1
    y_pad = (y_max - y_min) * 0.1
    x_min -= x_pad
    x_max += x_pad
    y_min -= y_pad
    y_max += y_pad
    
    # 调整x和y范围为等长，确保正方形显示
    dx = x_max - x_min
    dy = y_max - y_min
    d = max(dx, dy)
    x_center = (x_min + x_max) / 2
    y_center = (y_min + y_max) / 2
    x_min = x_center - d / 2
    x_max = x_center + d / 2
    y_min = y_center - d / 2
    y_max = y_center + d / 2
    
    hist, _, _ = np.histogram2d(x_data, y_data, bins=40, range=[[x_min, x_max], [y_min, y_max]])
    hist = np.where(hist == 0, np.nan, hist)  # 无数据区域设为透明
    # 取非NaN值的95%分位数作为vmax，增强对比
    valid_hist = hist[~np.isnan(hist)]
    vmax_val = np.percentile(valid_hist, 95) if valid_hist.size > 0 else 0
    im = ax.imshow(
        hist.T,  # 转置hist
        extent=[x_min, x_max, y_min, y_max],
        origin='lower',
        cmap='viridis',  # 颜色映射
        aspect='equal',
        vmin=valid_hist.min() if valid_hist.size > 0 else 0,
        vmax=vmax_val
    )
    plt.colorbar(im, ax=ax, shrink=0.8, label='Charge Count')
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    ax.set_facecolor('white')  # 背景为白色

def plot_distributions(positions):
    if len(positions) == 0:
        print("错误：所有能量沉积事件生成0电荷（可能edep过小或单位错误）")
        return
    
    x, y, z = positions.T
    
    # 三维散点图
    fig = plt.figure(figsize=(12, 10))
    ax3d = fig.add_subplot(221, projection='3d')
    ax3d.scatter(x, y, z, c='blue', alpha=0.3, s=5, label='Charge Points')
    ax3d.set_xlabel('X (mm)')
    ax3d.set_ylabel('Y (mm)')
    ax3d.set_zlabel('Z (mm)')
    ax3d.set_title('3D Charge Distribution with Diffusion')
    ax3d.set_facecolor('white')  # 背景为白色
    
    # 平面投影直方图
    ax_xy = fig.add_subplot(222)
    plot_2d_projection(ax_xy, x, y, 'XY Plane Projection', 'X (mm)', 'Y (mm)')
    
    ax_yz = fig.add_subplot(223)
    plot_2d_projection(ax_yz, y, z, 'YZ Plane Projection', 'Y (mm)', 'Z (mm)')
    
    ax_xz = fig.add_subplot(224)
    plot_2d_projection(ax_xz, x, z, 'XZ Plane Projection', 'X (mm)', 'Z (mm)')
    
    plt.tight_layout()
    plt.savefig('charge_distribution.png', dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    csv_path = "/opt/star-xp-master/build/source/0-Simulation/e-+,0.01,10/event_0.csv"
    deposits = read_deposits_info(csv_path)
    
    if not deposits:
        print("错误：未读取到任何能量沉积信息")
        exit(1)
    
    charge_positions = process_deposits(deposits)
    plot_distributions(charge_positions)
